# -*- coding: utf8 -*-

import sys
from collections import defaultdict

from Bio import SeqIO

import config
from config import AlignDatabaseLine, ExonInfo
from utils import profile, deprecated, dump_object, load_object


ALIGN_RESULTS_FILENAME = config.ALIGN_RESULTS_FILENAME
SCAFFOLDS_FILENAME = config.SCAFFOLDS_FILENAME
ALIGN_DUMP_FILENAME = config.ALIGN_DUMP_FILENAME
SEGMENT_LENGTH_DUMP_FILENAME = config.SEGMENT_LENGTH_DUMP_FILENAME
EXON_ALIGN_DUMP_FILENAME = config.EXON_ALIGN_DUMP_FILENAME
STRUCTURE_FILENAME = config.STRUCTURE_FILENAME


def main():
    print('Align results dump')
    make_align_results_dump()
    print('Segment length dump')
    make_segment_length_dump()
    print('Exon align dump')
    make_exon_align_dump()


@profile
def make_align_results_dump(dump_filename=ALIGN_DUMP_FILENAME, align_results_filename=ALIGN_RESULTS_FILENAME):
    """ segment_name -> [((gene_name, exon_name, exon_start, exon_end, strand),
                          query_name, segment_name, identical_match_percentage, alignment_length,
                          mismatch_count, gap_open_count, query_start, query_end, segment_start, segment_end,
                          e_value, bit_score)]
    """
    sys.stderr.write('Load align result...\n')
    segments = defaultdict(list)
    with open(align_results_filename) as handle:
        for line in handle.readlines():
            line_data = line.split()
            align = [tuple(line_data[0].split('|'))]
            # noinspection PyTypeChecker
            align.extend(line_data)
            line_data = tuple(align)
            segments[line_data[2]].append(line_data)
    sys.stderr.write('Dump align result...\n')
    dump_object(segments, dump_filename)
    sys.stderr.write('Done.\n')
    return segments


@profile
def make_segment_length_dump(dump_filename=SEGMENT_LENGTH_DUMP_FILENAME, scaffolds_filename=SCAFFOLDS_FILENAME):
    """ segment_name -> segment_length """
    result = {}
    with open(scaffolds_filename) as handle:
        for record in SeqIO.parse(handle, 'fasta'):
            result[record.id] = len(record.seq)
    dump_object(result, dump_filename)


@profile
def make_exon_align_dump(dump_filename=EXON_ALIGN_DUMP_FILENAME, align_dump_filename=ALIGN_DUMP_FILENAME):
    """ exon_name -> {AlignDatabaseLine :: (segment_id, align_start, strand)}"""
    sys.stderr.write('Load align...\n')
    aligns = load_object(align_dump_filename)
    new_data = defaultdict(set)
    for segment_aligns in aligns.values():
        for align in segment_aligns:
            if not check_align_line(align):
                continue
            start = int(align[9])
            end = int(align[10])
            exon_name = align[0][1]
            gene_strand = (end - start) // abs(end-start)
            align_db_line = AlignDatabaseLine(segment_id=int(align[2][4:]),
                                              start=min(start, end), end=max(start, end), strand=gene_strand)
            new_data[exon_name].add(align_db_line)
    dump_object(new_data, dump_filename)


def check_align_line(align_line):
    """ quality control for AlignDatabaseLine """
    align_length = int(align_line[4]) - int(align_line[5])
    exon_length = int(align_line[0][3]) + 1 - int(align_line[0][2])
    if (exon_length - align_length) / exon_length > config.EXON_PERCENT_MATCH_SCORE:
        return False
    align_e_value = float(align_line[11])
    if align_e_value > 1e-8:
        return False
    return True


@deprecated
@profile
def __parse_exon_length(structure_filename=STRUCTURE_FILENAME):
    """ exon_name -> (length, start, position_in_gene)"""
    sys.stderr.write('Parse exon length...\n')
    genes = defaultdict(set)
    with open(structure_filename) as handle:
        handle.readline()
        for line in handle.readlines():
            data = line.split()
            gene = data[0]
            name = data[1]
            start = int(data[2])
            end = int(data[3])
            genes[gene].add((start, end, name))
    results = defaultdict(int)
    for gene, exons in genes.items():
        exons = list(sorted(exons))
        for position, info in enumerate(exons):
            results[info[2]] = ExonInfo(length=info[1] + 1 - info[0], start=info[0], position_in_gene=position)
    sys.stderr.write('Done.\n')
    return results


if __name__ == '__main__':
    main()
